package coordinator

import (
	"context"
	"distribute-mapreduce/grpc/generated/pb/coordinator"
	"fmt"
	"github.com/sirupsen/logrus"
	"google.golang.org/grpc"
	"log"
	"net"
	"sync"
	"time"
)

type statusValue struct {
	status int // 0表示未分发，1表示已分发, 2表示完成
	id     int
	timer  *time.Timer
}

type Coordinator struct {
	// Your definitions here.
	coordinator.UnimplementedCoordinatorServiceServer
	fileStatus       map[string]*statusValue // 文件名-> 是否完成map阶段
	fileStatusReduce map[int]*statusValue    // 文件名-> 是否完成reduce阶段

	lock sync.RWMutex

	lockReduce sync.RWMutex

	nReduce int

	remainReduceId int // 还剩下多少个reduce任务没有完成

	port int

	listener net.Listener

	/*
		我们需要在第一个worker请求时开始计时，在所有worker完成任务后停止计时
	*/

	// 任务计时需要的类型
	timerStarter sync.Once
	timerStopper  sync.Once
	startTime time.Time
	stopTime time.Time
}

func (c *Coordinator) GetTask(ctx context.Context, param *coordinator.GetTaskParam) (*coordinator.GetTaskReply, error) {
	c.lock.Lock()
	defer c.lock.Unlock()
	// 1. 遍历文件列表，找到一个没有完成的文件
	c.timerStarter.Do(func() {
		// 顺便计个时
		c.startTime = time.Now()
	})
	for filename, status := range c.fileStatus {
		curStatus := status
		if curStatus.status == 0 {
			// 2. 找到一个没有完成的文件，返回给worker
			reply := &coordinator.GetTaskReply{
				Filename:  filename,
				TaskType:  coordinator.MRTaskType_MAP,
				MapTaskId: int32(status.id), // 一个文件对应一个map任务
				NReduce:   int32(c.nReduce),
			}
			curStatus.status = 1 // 标记为已分发

			// 如果5秒没有完成，就把状态改为未分发
			timer := time.AfterFunc(5*time.Second, func() {
				// 检查一下状态是不是还是1，如果是的话，就把它重新变成0
				c.lock.Lock()
				defer c.lock.Unlock()
				if curStatus.status == 1 {
					curStatus.status = 0 // 变成未分发
				}
			})
			curStatus.timer = timer
			return reply, nil
		}
	}

	// 1.1 再次检查是否有map任务处于已分发状态
	for _, status := range c.fileStatus {
		if status.status == 1 {
			return &coordinator.GetTaskReply{
				TaskType: coordinator.MRTaskType_WAIT,
			}, nil
		}
	}

	// 2. map任务均完成了，通知reduce任务
	c.lockReduce.Lock()
	defer c.lockReduce.Unlock()
	for reduceId, status := range c.fileStatusReduce {
		curStatus := status
		if curStatus.status == 0 {
			reply := &coordinator.GetTaskReply{
				TaskType: coordinator.MRTaskType_REDUCE,
				ReduceId: int32(reduceId),
				NReduce:  int32(c.nReduce),
			}
			curStatus.status = 1 // 标记为已分发
			// 如果5秒没有完成，就把状态改为未分发
			timer := time.AfterFunc(5*time.Second, func() {
				// 检查一下状态是不是还是1，如果是的话，就把它重新变成0
				c.lockReduce.Lock()
				defer c.lockReduce.Unlock()
				if curStatus.status == 1 {
					curStatus.status = 0 // 变成未分发
				}
			})
			curStatus.timer = timer
			logrus.Infof("[coordinator]-分发reduceid = %d的任务", reduceId)
			return reply, nil
		}
	}
	// 2.1 检查是否有reduce任务处于已分发状态
	for _, status := range c.fileStatusReduce {
		if status.status == 1 {
			return &coordinator.GetTaskReply{
				TaskType: coordinator.MRTaskType_WAIT,
			}, nil
		}
	}

	// 3. 所有任务都完成了，通知worker退出
	c.timerStopper.Do(func() {
		c.stopTime = time.Now()
	})
	runTime := c.stopTime.UnixMicro() - c.startTime.UnixMicro()
	logrus.Infof("[coordinator]-所有任务完成，运行时间为: %dms", runTime)
	defer c.listener.Close() // 完成任务后关闭grpc监听
	return &coordinator.GetTaskReply{
		TaskType: coordinator.MRTaskType_EXIT,
	}, nil
}

func (c *Coordinator) FinishMapTask(ctx context.Context, param *coordinator.FinishMapTaskParam) (*coordinator.FinishMapTaskReply, error) {
	c.lock.Lock()
	defer c.lock.Unlock()
	if _, ok := c.fileStatus[param.Filename]; ok {
		c.fileStatus[param.Filename].status = 2 // 标记为完成
		if c.fileStatus[param.Filename].timer != nil {
			c.fileStatus[param.Filename].timer.Stop() // 停止定时器
			c.fileStatus[param.Filename].timer = nil  // help gc
		}
	}
	return &coordinator.FinishMapTaskReply{
		OK: true,
	}, nil
}

func (c *Coordinator) FinishReduceTask(ctx context.Context, param *coordinator.FinishReduceTaskParam) (*coordinator.FinishReduceTaskReply, error) {
	c.lockReduce.Lock()
	defer c.lockReduce.Unlock()
	if _, ok := c.fileStatusReduce[int(param.ReduceId)]; ok {
		c.fileStatusReduce[int(param.ReduceId)].status = 2 // 标记为完成
		if c.fileStatusReduce[int(param.ReduceId)].timer != nil {
			c.fileStatusReduce[int(param.ReduceId)].timer.Stop() // 停止定时器
			c.fileStatusReduce[int(param.ReduceId)].timer = nil  // help gc
		}
	}
	return &coordinator.FinishReduceTaskReply{
		OK: true,
	}, nil
}
func (c *Coordinator) Run() {
	// 启动grpc服务
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", c.port))
	c.listener = listener
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	s := grpc.NewServer()
	coordinator.RegisterCoordinatorServiceServer(s, c)
	if err := s.Serve(listener); err != nil {
		log.Fatalf("failed to serve: %v", err)
	}
}

func (c *Coordinator) Hello(ctx context.Context, param *coordinator.HelloParam) (*coordinator.HelloReply, error) {
	return &coordinator.HelloReply{
		Msg: "hello",
	}, nil
}

func NewCoordinator(files []string, nReduce int, port int) *Coordinator {
	c := &Coordinator{
		fileStatus:       make(map[string]*statusValue),
		fileStatusReduce: make(map[int]*statusValue),
		nReduce:          nReduce,
		remainReduceId:   nReduce,
		port:             port,
	}
	fileId := 0
	for _, fileName := range files {
		c.fileStatus[fileName] = &statusValue{
			status: 0,
			id:     fileId,
		}
		fileId++
	}

	for i := 0; i < nReduce; i++ {
		c.fileStatusReduce[i] = &statusValue{
			status: 0,
			id:     i,
		}
	}
	return c
}
